Training a pistol detector with YOLOv5 and Neuralmagic

• Step 1: Download the pistol dataset
• Step 2: Explore the data
• Step 3: Split the dataset into train/test/val subsets
• Step 4: Install DeepSparse and SparseML libraries and train our model
• Step 4: Export the Sparse Model to ONNX format for CPU inference
• Step 5: DeepSparse Engine model deployment
• Step 6: Compare sparsified ONNX model with PyTorch version
• Step 7: Run the model on a YouTube video using Ultralytics library

Note: To run this notebook yourself please click on the "Open In Colab" tag above.

Step 1: Download the pistol dataset

Download the Pistols dataset from Roboflow's public database.

Select YOLOv5 from the download formats and copy your unique terminal command into the cell below.

! curl -L "https://public.roboflow.com/ds/xXg040ThIk?key=[YOUR_UNIQUE_KEY]" > data.zip \
  && unzip -q data.zip -d downloads \
  && rm data.zip

  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
100   901  100   901    0     0   3575      0 --:--:-- --:--:-- --:--:--  3575
100 58.4M  100 58.4M    0     0  70.8M      0 --:--:-- --:--:-- --:--:-- 70.8M

Step 2: Explore the data

To explore the rapid dataset exploration toolset from Fiftyone.

To use Fiftyone with Colab we need to install opencv-python-headless==4.5.4.60.

! pip install -q opencv-python-headless==4.5.4.60 fiftyone

import fiftyone as fo

# Import dataset by explicitly providing paths to the source images and labels
dataset = fo.Dataset.from_dir(
    dataset_type=fo.types.YOLOv4Dataset,
    data_path="downloads/export/images",
    label_field="ground_truth",
    labels_path="downloads/export/labels",
    classes=["pistol"],
)

# View the data using the interactive dashboard
session = fo.launch_app(dataset)

The pistol images seem to be correctly labelled with a good variety of different pistols included in the dataset. The image size is 416x416 for all images.

Step 3: Split the dataset into train/test/val subsets

Next we can use Fiftyone utilities to split the dataset into three subsets. We define a convenience function called "split_dataset" to help us process the data.

import fiftyone.utils.splits as fous
import os


def split_dataset(dataset, split_dict, output_path):
    fous.random_split(dataset, split_dict, seed=42)
    for tag in list(split_dict.keys()):
        view = dataset.match_tags(tag)
        # Export each split separately
        dir = os.path.join(output_path, tag)
        view.export(
            export_dir=dir,
            dataset_type=fo.types.YOLOv4Dataset,
            label_field="ground_truth",
        )


split_dataset(
    dataset=dataset,
    split_dict={"train": 0.8, "test": 0.1, "val": 0.1},
    output_path="dataset/guns/",
)

 100% |███████████████| 2377/2377 [2.7s elapsed, 0s remaining, 859.1 samples/s]      
 100% |█████████████████| 297/297 [355.2ms elapsed, 0s remaining, 836.1 samples/s]      
 100% |█████████████████| 297/297 [397.3ms elapsed, 0s remaining, 747.5 samples/s]      

To use the dataset with a YOLO training pipeline we need to define a *.yml file that points to the data locations. To read more about training on a custom dataset using YOLOv5, see Train Custom Data by the Ultralytics team.

content = """# Train/val/test sets
path: /content/dataset/guns  # dataset root dir
train: train/data  # train images (relative to 'path')
test: test/data  # test images (relative to 'path')
val: val/data  # val images (relative to 'path')

# classes
nc: 1 # number of classes
names: ['pistol']"""

with open("pistols.yml", "w") as f:
    f.write(content)

Remove the downloaded data to clean up the workspace.

! rm -r downloads

Step 4: Install DeepSparse and SparseML libraries and train our model

To train our own custom model using the pistol dataset we will use the SparseML YOLOv5 Integration provided by Neuralmagic.

First, we install the required libraries and we grab the transfer learning recipe from their repo.

Note, you can also use the built-in recipe, like this:

--recipe zoo:cv/detection/yolov5-s/pytorch/ultralytics/coco/pruned-aggressive_96

However, if you want to specify the number of epochs or the learning rate rather use the downloaded recipe.

! pip install -q deepsparse[yolo] sparseml[torchvision]

! curl -LO "https://raw.githubusercontent.com/neuralmagic/sparseml/main/integrations/ultralytics-yolov5/recipes/yolov5.transfer_learn_pruned.md"

Now we can train our model using transfer learning on YOLOv5s (small) and the downloaded recipe. We use the CLI and run sparseml.yolov5.train

The arguments we use with sparseml.yolov5.train are shown below. For more info on each argument, just run:

! sparseml.yolov5.train --help

For more information on the pre-trained weights you can use with YOLOv5 feel free to refer to the Sparse Transfer Learning With YOLOv5 example.

We train our model for 50 epochs. Click on "Show Output" to view the results

! sparseml.yolov5.train \
  --data pistols.yml \
  --cfg models_v5.0/yolov5s.yaml \
  --weights zoo:cv/detection/yolov5-s/pytorch/ultralytics/coco/pruned-aggressive_96 \
  --hyp data/hyps/hyp.finetune.yaml \
  --batch-size -1 \
  --imgsz 416 \
  --recipe yolov5.transfer_learn_pruned.md

train: weights=zoo:cv/detection/yolov5-s/pytorch/ultralytics/coco/pruned-aggressive_96, cfg=models_v5.0/yolov5s.yaml, data=pistols.yml, hyp=data/hyps/hyp.finetune.yaml, epochs=300, batch_size=-1, imgsz=416, rect=False, resume=False, nosave=False, noval=False, noautoanchor=False, evolve=None, bucket=, cache=None, image_weights=False, device=, multi_scale=False, single_cls=False, optimizer=SGD, sync_bn=False, workers=8, project=yolov5_runs/train, name=exp, exist_ok=False, quad=False, cos_lr=False, label_smoothing=0.0, patience=0, freeze=[0], save_period=-1, local_rank=-1, entity=None, upload_dataset=False, bbox_interval=-1, artifact_alias=latest, recipe=yolov5.transfer_learn_pruned.md, disable_ema=False, max_train_steps=-1, max_eval_steps=-1, one_shot=False, num_export_samples=0
github: skipping check (not a git repository), for updates see https://github.com/ultralytics/yolov5
fatal: not a git repository (or any of the parent directories): .git
YOLOv5 🚀 2022-6-30 torch 1.9.1+cu102 CUDA:0 (Tesla T4, 15110MiB)

hyperparameters: lr0=0.0032, lrf=0.12, momentum=0.843, weight_decay=0.00036, warmup_epochs=2.0, warmup_momentum=0.5, warmup_bias_lr=0.05, box=0.0296, cls=0.243, cls_pw=0.631, obj=0.301, obj_pw=0.911, iou_t=0.2, anchor_t=2.91, fl_gamma=0.0, hsv_h=0.0138, hsv_s=0.664, hsv_v=0.464, degrees=0.373, translate=0.245, scale=0.898, shear=0.602, perspective=0.0, flipud=0.00856, fliplr=0.5, mosaic=1.0, mixup=0.243, copy_paste=0.0
Weights & Biases: run 'pip install wandb' to automatically track and visualize YOLOv5 🚀 runs (RECOMMENDED)
TensorBoard: Start with 'tensorboard --logdir yolov5_runs/train', view at http://localhost:6006/
Obtaining new sparse zoo credentials token
Getting signed url for c13e55cb-dd6c-4492-a079-8986af0b65e6/model.pt
Downloading model file model.pt to /root/.cache/sparsezoo/c13e55cb-dd6c-4492-a079-8986af0b65e6/pytorch/model.pt
downloading...: 100% 14.1M/14.1M [00:00<00:00, 37.0MB/s]
Overriding model.yaml nc=80 with nc=1

                 from  n    params  module                                  arguments                     
  0                -1  1      3520  yolov5.models.common.Focus              [3, 32, 3]                    
  1                -1  1     18560  yolov5.models.common.Conv               [32, 64, 3, 2]                
  2                -1  1     18816  yolov5.models.common.C3                 [64, 64, 1]                   
  3                -1  1     73984  yolov5.models.common.Conv               [64, 128, 3, 2]               
  4                -1  3    156928  yolov5.models.common.C3                 [128, 128, 3]                 
  5                -1  1    295424  yolov5.models.common.Conv               [128, 256, 3, 2]              
  6                -1  3    625152  yolov5.models.common.C3                 [256, 256, 3]                 
  7                -1  1   1180672  yolov5.models.common.Conv               [256, 512, 3, 2]              
  8                -1  1    656896  yolov5.models.common.SPP                [512, 512, [5, 9, 13]]        
  9                -1  1   1182720  yolov5.models.common.C3                 [512, 512, 1, False]          
 10                -1  1    131584  yolov5.models.common.Conv               [512, 256, 1, 1]              
 11                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
 12           [-1, 6]  1         0  yolov5.models.common.Concat             [1]                           
 13                -1  1    361984  yolov5.models.common.C3                 [512, 256, 1, False]          
 14                -1  1     33024  yolov5.models.common.Conv               [256, 128, 1, 1]              
 15                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
 16           [-1, 4]  1         0  yolov5.models.common.Concat             [1]                           
 17                -1  1     90880  yolov5.models.common.C3                 [256, 128, 1, False]          
 18                -1  1    147712  yolov5.models.common.Conv               [128, 128, 3, 2]              
 19          [-1, 14]  1         0  yolov5.models.common.Concat             [1]                           
 20                -1  1    296448  yolov5.models.common.C3                 [256, 256, 1, False]          
 21                -1  1    590336  yolov5.models.common.Conv               [256, 256, 3, 2]              
 22          [-1, 10]  1         0  yolov5.models.common.Concat             [1]                           
 23                -1  1   1182720  yolov5.models.common.C3                 [512, 512, 1, False]          
 24      [17, 20, 23]  1     16182  yolov5.models.yolo.Detect               [1, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]]
overriding activations in model to Hardswish
YOLOv5s summary: 283 layers, 7063542 parameters, 7063542 gradients, 16.5 GFLOPs

Transferred 354/361 items from /root/.cache/sparsezoo/c13e55cb-dd6c-4492-a079-8986af0b65e6/pytorch/model.pt
AutoBatch: Computing optimal batch size for --imgsz 416
AutoBatch: CUDA:0 (Tesla T4) 14.76G total, 0.06G reserved, 0.05G allocated, 14.64G free
      Params      GFLOPs  GPU_mem (GB)  forward (ms) backward (ms)                   input                  output
     7063542       6.961         0.201         16.78         14.15        (1, 3, 416, 416)                    list
     7063542       13.92         0.310         17.35         15.45        (2, 3, 416, 416)                    list
     7063542       27.85         0.581         18.58         18.79        (4, 3, 416, 416)                    list
     7063542       55.69         1.126         19.02         26.23        (8, 3, 416, 416)                    list
     7063542       111.4         2.223         32.87         50.45       (16, 3, 416, 416)                    list
AutoBatch: Using batch-size 96 for CUDA:0 13.28G/14.76G (90%)
Scaled weight_decay = 0.00054
optimizer: SGD with parameter groups 59 weight (no decay), 62 weight, 62 bias
albumentations: version 1.0.3 required by YOLOv5, but version 0.1.12 is currently installed
train: Scanning '/content/dataset/guns/train/data' images and labels...2377 found, 0 missing, 0 empty, 0 corrupt: 100% 2377/2377 [00:01<00:00, 1379.73it/s]
train: New cache created: /content/dataset/guns/train/data.cache
val: Scanning '/content/dataset/guns/val/data' images and labels...297 found, 0 missing, 0 empty, 0 corrupt: 100% 297/297 [00:00<00:00, 634.83it/s]
val: New cache created: /content/dataset/guns/val/data.cache
Plotting labels to yolov5_runs/train/exp2/labels.jpg... 

AutoAnchor: 2.95 anchors/target, 1.000 Best Possible Recall (BPR). Current anchors are a good fit to dataset ✅
Image sizes 416 train, 416 val
Using 2 dataloader workers
Logging results to yolov5_runs/train/exp2
Starting training for 300 epochs...
Disabling LR scheduler, managing LR using SparseML recipe
Overriding number of epochs from SparseML manager to 50

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      0/49     12.2G   0.07143  0.005311         0       243       416: 100% 25/25 [00:35<00:00,  1.43s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:04<00:00,  2.16s/it]
                 all        297        357     0.0179      0.185     0.0106    0.00189

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      1/49     14.3G    0.0591  0.005631         0       237       416: 100% 25/25 [00:32<00:00,  1.30s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.95s/it]
                 all        297        357       0.11      0.199     0.0727     0.0206

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      2/49     12.3G    0.0511  0.005889         0       186       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.61s/it]
                 all        297        357      0.356      0.401      0.275     0.0831

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      3/49     12.9G   0.04455  0.006083         0       195       416: 100% 25/25 [00:33<00:00,  1.33s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.53s/it]
                 all        297        357      0.419      0.473      0.406      0.132

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      4/49     12.9G   0.03806  0.006098         0       233       416: 100% 25/25 [00:32<00:00,  1.32s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.50s/it]
                 all        297        357      0.555      0.605      0.598      0.259

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      5/49     12.9G   0.03427  0.005962         0       258       416: 100% 25/25 [00:32<00:00,  1.30s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.46s/it]
                 all        297        357      0.729      0.591      0.686      0.347

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      6/49     12.9G   0.03197   0.00576         0       207       416: 100% 25/25 [00:33<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.47s/it]
                 all        297        357      0.718      0.669      0.718      0.392

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      7/49     12.9G   0.03004  0.005625         0       254       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.52s/it]
                 all        297        357      0.749      0.669      0.737      0.411

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      8/49       13G   0.02922  0.005537         0       235       416: 100% 25/25 [00:34<00:00,  1.38s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.44s/it]
                 all        297        357      0.701      0.691      0.752      0.436

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      9/49       13G   0.02814  0.005131         0       196       416: 100% 25/25 [00:33<00:00,  1.35s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.47s/it]
                 all        297        357      0.717      0.709      0.772      0.447

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     10/49       13G   0.02756  0.005152         0       247       416: 100% 25/25 [00:34<00:00,  1.37s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.48s/it]
                 all        297        357      0.671      0.787       0.78      0.461

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     11/49       13G   0.02719  0.005132         0       227       416: 100% 25/25 [00:32<00:00,  1.30s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.51s/it]
                 all        297        357      0.847      0.666        0.8      0.463

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     12/49       13G   0.02633  0.004897         0       265       416: 100% 25/25 [00:32<00:00,  1.32s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.48s/it]
                 all        297        357      0.722      0.773      0.811      0.495

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     13/49       13G   0.02559  0.005043         0       196       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.41s/it]
                 all        297        357      0.691      0.814      0.791      0.469

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     14/49       13G   0.02537  0.004815         0       232       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.44s/it]
                 all        297        357      0.791      0.731      0.811      0.499

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     15/49       13G   0.02512  0.004849         0       230       416: 100% 25/25 [00:34<00:00,  1.37s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.45s/it]
                 all        297        357      0.765      0.739       0.81      0.483

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     16/49       13G   0.02504  0.004759         0       206       416: 100% 25/25 [00:33<00:00,  1.33s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.47s/it]
                 all        297        357      0.737      0.779       0.82      0.512

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     17/49       13G   0.02469  0.004673         0       267       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.41s/it]
                 all        297        357      0.802       0.72      0.822      0.512

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     18/49       13G   0.02447  0.004662         0       231       416: 100% 25/25 [00:33<00:00,  1.35s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.46s/it]
                 all        297        357      0.778      0.759      0.824      0.522

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     19/49       13G   0.02424  0.004754         0       235       416: 100% 25/25 [00:33<00:00,  1.35s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.39s/it]
                 all        297        357      0.785      0.759      0.827      0.533

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     20/49       13G   0.02374  0.004643         0       213       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.42s/it]
                 all        297        357      0.857      0.695      0.823      0.517

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     21/49       13G   0.02402  0.004533         0       209       416: 100% 25/25 [00:34<00:00,  1.37s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.46s/it]
                 all        297        357      0.796      0.754      0.816      0.523

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     22/49     13.6G    0.0235  0.004528         0       236       416: 100% 25/25 [00:32<00:00,  1.31s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.40s/it]
                 all        297        357      0.777      0.782      0.833      0.545

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     23/49     13.6G   0.02366  0.004616         0       264       416: 100% 25/25 [00:33<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.44s/it]
                 all        297        357      0.816      0.748      0.832      0.531

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     24/49     13.6G   0.02322  0.004491         0       238       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.49s/it]
                 all        297        357      0.814      0.748      0.829      0.519

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     25/49     13.6G   0.02305  0.004363         0       213       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.40s/it]
                 all        297        357      0.768      0.773      0.835      0.536

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     26/49     13.6G   0.02326  0.004525         0       238       416: 100% 25/25 [00:33<00:00,  1.33s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.43s/it]
                 all        297        357      0.801      0.787      0.849      0.552

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     27/49     13.6G   0.02263  0.004288         0       250       416: 100% 25/25 [00:33<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.38s/it]
                 all        297        357      0.813       0.77      0.845      0.548

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     28/49     13.6G   0.02294  0.004502         0       245       416: 100% 25/25 [00:33<00:00,  1.33s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.46s/it]
                 all        297        357      0.821      0.762       0.84      0.546

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     29/49     13.6G   0.02249  0.004358         0       223       416: 100% 25/25 [00:33<00:00,  1.35s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.42s/it]
                 all        297        357      0.811       0.78      0.839      0.547

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     30/49     13.7G    0.0226  0.004365         0       234       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.40s/it]
                 all        297        357      0.832       0.75      0.837      0.545

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     31/49     13.7G   0.02216  0.004202         0       218       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.39s/it]
                 all        297        357      0.847      0.744      0.841      0.542

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     32/49     13.7G   0.02241  0.004421         0       258       416: 100% 25/25 [00:33<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.45s/it]
                 all        297        357      0.789      0.798      0.841      0.553

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     33/49     13.7G   0.02187  0.004324         0       183       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.42s/it]
                 all        297        357      0.798      0.798      0.847      0.551

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     34/49     13.7G   0.02235  0.004243         0       217       416: 100% 25/25 [00:32<00:00,  1.30s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.45s/it]
                 all        297        357      0.821      0.773      0.843       0.55

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     35/49     13.7G   0.02182  0.004277         0       234       416: 100% 25/25 [00:33<00:00,  1.35s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.42s/it]
                 all        297        357      0.804      0.796      0.849      0.559

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     36/49     13.7G   0.02203  0.004312         0       260       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.47s/it]
                 all        297        357      0.809      0.784      0.846      0.553

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     37/49     13.7G   0.02222   0.00428         0       200       416: 100% 25/25 [00:32<00:00,  1.31s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.43s/it]
                 all        297        357      0.806      0.804       0.85      0.557

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     38/49     13.7G   0.02202  0.004223         0       210       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.46s/it]
                 all        297        357      0.784      0.796      0.836      0.551

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     39/49     13.7G   0.02193  0.004199         0       256       416: 100% 25/25 [00:34<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.43s/it]
                 all        297        357      0.804      0.795      0.844      0.556

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     40/49     13.7G   0.02192  0.004291         0       246       416: 100% 25/25 [00:34<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.24s/it]
                 all        297        357      0.792        0.8      0.839      0.558

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     41/49     13.7G   0.02164  0.004269         0       235       416: 100% 25/25 [00:32<00:00,  1.31s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.47s/it]
                 all        297        357      0.783       0.81      0.842      0.561

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     42/49     13.7G   0.02207  0.004228         0       247       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.43s/it]
                 all        297        357       0.79      0.812      0.845      0.562

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     43/49     13.7G   0.02193  0.004293         0       263       416: 100% 25/25 [00:33<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.57s/it]
                 all        297        357      0.796      0.801       0.85      0.561

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     44/49     13.7G   0.02202  0.004254         0       242       416: 100% 25/25 [00:34<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.44s/it]
                 all        297        357       0.81      0.793      0.849      0.556

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     45/49     13.7G   0.02163  0.004189         0       209       416: 100% 25/25 [00:33<00:00,  1.33s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.48s/it]
                 all        297        357      0.791      0.807      0.849      0.557

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     46/49     13.7G   0.02168  0.004311         0       241       416: 100% 25/25 [00:34<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.44s/it]
                 all        297        357      0.801      0.801      0.844      0.556

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     47/49     13.7G   0.02202  0.004218         0       231       416: 100% 25/25 [00:33<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.44s/it]
                 all        297        357      0.811      0.793      0.847      0.559

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     48/49     13.7G   0.02156  0.004261         0       230       416: 100% 25/25 [00:32<00:00,  1.32s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.51s/it]
                 all        297        357      0.835      0.765      0.842      0.556

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     49/49     13.7G   0.02169  0.004359         0       239       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.43s/it]
                 all        297        357      0.803       0.81      0.851      0.561

51 epochs completed in 0.526 hours.
Optimizer stripped from yolov5_runs/train/exp2/weights/last.pt, 14.4MB
Optimizer stripped from yolov5_runs/train/exp2/weights/best.pt, 14.4MB

Validating yolov5_runs/train/exp2/weights/best.pt...
Fusing layers... 
YOLOv5s summary: 224 layers, 7053910 parameters, 0 gradients, 16.3 GFLOPs
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.97s/it]
                 all        297        357      0.789      0.815      0.845      0.562
Results saved to yolov5_runs/train/exp2

Step 4: Export the Sparse Model to ONNX format for CPU inference

The DeepSparse Engine accepts ONNX formats and is engineered to significantly speed up inference on CPUs for the sparsified models from this integration. To convert the YOLOv5 model to ONNX format run:

! sparseml.yolov5.export_onnx \
  --weights yolov5_runs/train/exp/weights/best.pt \
  --imgsz 416 \
  --dynamic

Step 5: DeepSparse Engine model deployment

To validate the model on our test set in the DeepSparse Engine we run the following command, again using the CLI interface.

Note that the DeepSparse engine uses the CPU for inference and it might take a minute or two to run through the test set.

! sparseml.yolov5.val_onnx \
  --model_path yolov5_runs/train/exp/weights/best.onnx \
  --batch-size 33 \
  --imgsz 416 \
  --task test \
  --data pistols.yml

val_onnx: data=pistols.yml, model_path=yolov5_runs/train/exp/weights/best.onnx, batch_size=33, imgsz=416, conf_thres=0.001, iou_thres=0.6, task=test, device=, workers=8, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=False, project=../usr/local/lib/python3.7/dist-packages/sparseml/yolov5/runs/val, name=exp, exist_ok=False, half=False, dnn=False, engine=deepsparse, num_cores=None
val_onnx: data=pistols.yml, model_path=yolov5_runs/train/exp/weights/best.onnx, batch_size=33, imgsz=416, conf_thres=0.001, iou_thres=0.6, task=test, device=, workers=8, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=False, project=../usr/local/lib/python3.7/dist-packages/sparseml/yolov5/runs/val, name=exp, exist_ok=False, half=False, dnn=False, engine=deepsparse, num_cores=None
fatal: not a git repository (or any of the parent directories): .git
YOLOv5 🚀 2022-6-30 torch 1.9.1+cu102 CUDA:0 (Tesla T4, 15110MiB)

DeepSparse Engine, Copyright 2021-present / Neuralmagic, Inc. version: 1.0.0 (8eaddc24) (release) (optimized) (system=avx2, binary=avx2)
test: Scanning '/content/dataset/guns/test/data' images and labels...297 found, 0 missing, 0 empty, 0 corrupt: 100% 297/297 [00:00<00:00, 1260.64it/s]
test: New cache created: /content/dataset/guns/test/data.cache
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 9/9 [00:34<00:00,  3.88s/it]
/usr/local/lib/python3.7/dist-packages/yolov5/utils/metrics.py:74: FutureWarning: elementwise comparison failed; returning scalar instead, but in the future will perform elementwise comparison
  names = [v for k, v in names.items() if k in unique_classes]  # list: only classes that have data
                 all        297        351      0.918      0.764      0.886      0.578
Speed: 0.0ms pre-process, 116.0ms inference, 0.0ms NMS per image at shape (33, 3, 416, 416)
Results saved to ../usr/local/lib/python3.7/dist-packages/sparseml/yolov5/runs/val/exp2

Step 6: Compare sparsified ONNX model with PyTorch version

We can also run the model trained with the Neuralmagic CLI using the Ultralytics YOLO library. We will compare the inference times for the ONNX model (*.onnx) run on the DeepSparse engine with the PyTorch version (*.pt) run on CPU.

We clone the YOLOv5 library and install dependencies.

! git clone https://github.com/ultralytics/yolov5
! pip install -q -r yolov5/requirements.txt

Now we run the ONNx model on a test image.

Then we run inference on the same image using the PyTorch model.

!python yolov5/detect.py \
  --weights /content/yolov5_runs/train/exp/weights/best.pt \
  --source  \
  --imgsz 416 \
  --device cpu

Step 7: Run the model on a YouTube video using Ultralytics library

Simply run inference with the YOLOv5 detect.py scrip using the .pt weights file instead of the .onnx one.

For this example we will use a YouTube scene from John Wick making Revolver to test the model. Inference will be done using the GPU this time.

!python yolov5/detect.py \
  --weights /content/yolov5_runs/train/exp/weights/best.pt \
  --source "https://www.youtube.com/watch?v=nB0A4pFIOA0" \
  --conf-thres=0.6 \
  --imgsz 416

Below are a few screenshots from the annotated video.